import argparse
import csv

def parse_args():
    parser = argparse.ArgumentParser(description="输出 tsv1 中同时在 tsv2 中的行到 tsv3")
    parser.add_argument("--tsv1", required=True, help="输入的 tsv1 文件路径")
    parser.add_argument("--tsv2", required=True, help="输入的 tsv2 文件路径")
    parser.add_argument("--tsv3", required=True, help="输出的 tsv3 文件路径")
    return parser.parse_args()

def read_tsv2_keys(tsv2_path):
    """
    读取 tsv2 文件中的键 (chrom, pos, alignstrand, readname)，并存储到集合中
    """
    keys = set()
    with open(tsv2_path, "r") as file:
        reader = csv.reader(file, delimiter="\t")
        for line in reader:
            # 构建键：chrom, pos, alignstrand, readname
            key = (line[0], line[1], line[2])
            keys.add(key)
    return keys

def filter_tsv1(tsv1_path, tsv2_keys, tsv3_path):
    """
    输出 tsv1 中在 tsv2 键集合中的行到 tsv3
    """
    with open(tsv1_path, "r") as infile, open(tsv3_path, "w", newline="") as outfile:
        reader = csv.reader(infile, delimiter="\t")
        writer = csv.writer(outfile, delimiter="\t")
        
        for line in reader:
            # 构建 tsv1 的键：chrom, pos, alignstrand, readname
            key = (line[0], line[1], line[2])
            if key in tsv2_keys:
                writer.writerow(line)

def main():
    args = parse_args()
    
    # 读取 tsv2 的键
    tsv2_keys = read_tsv2_keys(args.tsv2)
    
    # 过滤 tsv1 并输出到 tsv3
    filter_tsv1(args.tsv1, tsv2_keys, args.tsv3)

if __name__ == "__main__":
    main()
